#!/usr/bin/python3
# -*- coding:UTF-8 -*-
# Author:nigo

import os
from PIL import Image
import numpy as np
from wordcloud import WordCloud, ImageColorGenerator, get_single_color_func
import jieba
import sys
import random


def grey_color_func(word, font_size, position, orientation, random_state=None,
                    **kwargs):
    return "hsl(0, 0%%, %d%%)" % random.randint(60, 100)

def calc_word_frequency(text_path, stopwords_path, num=0):
    # 读取停用词表
    with open(stopwords_path, encoding='utf-8') as f_stop:
        f_stop_text = f_stop.read()
        f_stop_seg_list = f_stop_text.splitlines()
    #读入文件
    with open(text_path, encoding="utf-8") as fp:
        text = fp.read()
    words = jieba.lcut(text)#分词
    #统计词频
    counts= {}
    for word in words:
        if len(word) > 1 and word not in f_stop_seg_list:
            counts[word] = counts.get(word,0)+1
    #去标点（由于我这里不统计长度为1的词，去标点这步可省略）
    frequency = sorted(counts.items(), key=lambda item: item[1], reverse=True)
    if num > 0:
        frequency = frequency[:num]
        frequency_dict = {}
        for word in frequency:
            frequency_dict[word[0]] = word[1]
        return frequency_dict
    return counts

def main():
    # load image. This has been modified in gimp to be brighter and have more saturation.
    image_color = np.array(Image.open(image_path))
    wc = WordCloud(max_words=2000, 
                   mask = image_color,
                   random_state = 42, relative_scaling = 0, 
                   font_path = font_path,
                   width = 2560,
                   height = 1600,
                   background_color = 'white',
                   mode = 'RGBA',
                   )
    freqency = calc_word_frequency(text_path, stopwords_path, word_num)
    wc.generate_from_frequencies(freqency)
    if color_mode == 'color':
        image_colors = ImageColorGenerator(image_color)
        wc.recolor(color_func=image_colors)
    elif color_mode == 'grey':
        wc.recolor(color_func=grey_color_func, random_state=3)
    elif color_mode == 'default':
        pass
    else:
        wc.recolor(color_func=get_single_color_func(color_mode))
    wc.to_file(output_image)

def get_path(path):
    return os.path.join(d, path)

if len(sys.argv) < 5:
    print('''
    使用方法:
    wordcloud text_path image_path stopwords_path output_image color_mode word_num

    参数：
    color_mode:可以为'color', 'grey', 'default'
    word_num:为显示多少个词, 默认是全部
    ''')
else:
    d = os.getcwd()
    text_path = get_path(sys.argv[1])
    image_path = get_path(sys.argv[2])
    stopwords_path = get_path(sys.argv[3])
    output_image = get_path(sys.argv[4])
    font_path ='/usr/share/fonts/wenquanyi/wqy-microhei/wqy-microhei.ttc'
    if len(sys.argv) == 5:
        color_mode = 'default'
    else:
        color_mode = sys.argv[5]
        if len(sys.argv) == 6:
            word_num = 0
        else:
            word_num = int(sys.argv[6])
    main()
